import os
import argparse
from scripts.codeql_runner import run_callgraph_query
from scripts.neo4j_loader import get_driver, import_callgraph_csv


def cmd_run_query(args):
    run_callgraph_query(args.db, args.query, args.out)


def cmd_import(args):
    driver = get_driver()
    try:
        import_callgraph_csv(driver, args.csv)
    finally:
        driver.close()


def cmd_all(args):
    tmp_csv = args.out or os.path.join("out", "callgraph.csv")
    run_callgraph_query(args.db, args.query, tmp_csv)
    driver = get_driver()
    try:
        import_callgraph_csv(driver, tmp_csv)
    finally:
        driver.close()


def build_parser():
    parser = argparse.ArgumentParser(description="C++ 调用关系 → CodeQL → Neo4j 管线")
    sub = parser.add_subparsers(dest="cmd", required=True)

    p1 = sub.add_parser("run-query", help="运行 CodeQL 调用关系查询，导出 CSV")
    p1.add_argument("--db", required=True, help="CodeQL 数据库路径")
    p1.add_argument(
        "--query",
        default=os.path.join("queries", "CallGraph.ql"),
        help="查询文件路径 (默认 queries/CallGraph.ql)",
    )
    p1.add_argument("--out", required=True, help="输出 CSV 路径")
    p1.set_defaults(func=cmd_run_query)

    p2 = sub.add_parser("import-neo4j", help="将 CSV 导入 Neo4j")
    p2.add_argument("--csv", required=True, help="CodeQL 查询生成的 CSV 文件")
    p2.set_defaults(func=cmd_import)

    p3 = sub.add_parser("all", help="一键执行查询并导入 Neo4j")
    p3.add_argument("--db", required=True, help="CodeQL 数据库路径")
    p3.add_argument(
        "--query",
        default=os.path.join("queries", "CallGraph.ql"),
        help="查询文件路径 (默认 queries/CallGraph.ql)",
    )
    p3.add_argument("--out", help="输出 CSV 路径 (默认 out/callgraph.csv)")
    p3.set_defaults(func=cmd_all)

    return parser


if __name__ == "__main__":
    parser = build_parser()
    args = parser.parse_args()
    args.func(args)

